#!/usr/bin/env python
import pygame as pg
from random import randint,shuffle
from time import time
import sys
import threading
from math import floor


def clamp(x,lower,upper):
	if x > upper:
		x = upper
	elif x < lower:
		x = lower
	return x

class Polygon:
	points = []
	colour = (0,0,0)
	alpha = 255
	sf = None
	changed = True
	ymax = 0
	xmax = 0
	def __init__(self):
		pass
	def ReColour(self):
		self.colour = (randint(0,255),randint(0,255),randint(0,255))
		self.alpha = randint(0,255)
		self.changed = True
	def ReColourDelta(self,delta):
		newred = clamp(self.colour[0]+randint(-delta,delta),0,255)
		newgreen = clamp(self.colour[1]+randint(-delta,delta),0,255)
		newblue = clamp(self.colour[2]+randint(-delta,delta),0,255)
		self.colour = (newred,newgreen,newblue)
		self.alpha = clamp(self.alpha+randint(-delta,delta),0,255)
		self.changed = True
	def ReColourMedium(self):
		e = randint(0,4)
		(newred,newgreen,newblue) = self.colour
		if e == 0:
			newred = randint(0,255)
		elif e == 1:
			newgreen = randint(0,255)
		elif e == 2:
			newblue = randint(0,255)
		else:
			self.alpha = randint(0,255)
		if e < 4:
			self.colour = (newred,newgreen,newblue)
		self.changed = True
	def ReColourMediumDelta(self,delta):
		e = randint(0,4)
		(newred,newgreen,newblue) = self.colour
		if e == 0:
			newred = clamp(self.colour[0]+randint(-delta,delta),0,255)
		elif e == 1:
			newgreen = clamp(self.colour[1]+randint(-delta,delta),0,255)
		elif e == 2:
			newblue = clamp(self.colour[2]+randint(-delta,delta),0,255)
		else:
			self.alpha = clamp(self.alpha+randint(-delta,delta),0,255)
		if e < 4:
			self.colour = (newred,newgreen,newblue)
		self.changed = True
	def Generate(self,points,xmax,ymax):
		self.xmax = xmax
		self.ymax = ymax
		self.sf = pg.Surface((xmax,ymax))
		self.points = [0]*points
		for i in range(0,points):
			x,y = randint(0,xmax),randint(0,ymax)
			self.points[i] = (x,y)
		self.ReColour()

	def ReShape(self,delta):
		choice = randint(0,len(self.points)-1)
		x,y = randint(-delta,+delta),randint(-delta,+delta)
		x = clamp(x+self.points[choice][0],0,self.xmax)
		y = clamp(y+self.points[choice][1],0,self.ymax)
		self.points[choice] = (x,y)
		self.changed = True

	def Draw(self,surface):
		if self.changed:
			self.sf.fill((0,0,0))
			pg.draw.polygon(self.sf, self.colour, self.points)
			self.sf.set_colorkey((0,0,0))
			self.sf.set_alpha(self.alpha)
		surface.blit(self.sf,(0,0))

	def Clone(self):
		e = Polygon()
		e.points = self.points[0:]
		e.colour = self.colour
		e.alpha = self.alpha
		e.sf = self.sf
		e.changed = False
		e.xmax = self.xmax
		e.ymax = self.ymax
		return e


class Drawing:
	width = 0
	height = 0
	pointsPerPolygon = 6
	polygons = []
	surface = None
	length = None
	surfacedata = None
	def __init__(self,width,height):
		self.width = width
		self.height = height
	def Mutate(self, numMutations):
		for i in range(0,numMutations):
			e = randint(0,2)
			if e == 0:
				choice = randint(0,self.length-1)
				self.polygons[choice].ReColourMediumDelta(5)
			elif e == 1:
				choice = randint(0,self.length-1)
				self.polygons[choice].ReShape(25)
			elif e == 2:
				choice = randint(0,self.length-1)
				foo = randint(0,1)
				if foo == 0:
					next = (choice-1)%self.length
				else:
					next = (choice + 1)%self.length
				self.polygons[choice],self.polygons[next] = self.polygons[next],self.polygons[choice]
				#shuffle(self.polygons)
	def Draw(self,surface):
		assert surface.get_size() == (self.width,self.height)
		surface.fill((0,0,0))
		for polygon in self.polygons:
			polygon.Draw(surface)
	def Generate(self,polygons,sides):
		self.length = polygons
		self.polygons = [0]*polygons
		self.pointsPerPolygon = sides
		for i in range(0,polygons):
			e = Polygon()
			#e.ReColour()
			e.Generate(self.pointsPerPolygon,self.width,self.height)
			self.polygons[i] = e

	def Clone(self):
		d = Drawing(self.width,self.height)
		length = len(self.polygons)
		d.polygons = [0]*length
		for i in range(0,length):
			d.polygons[i] = self.polygons[i].Clone()
		d.length = self.length
		return d

#wrapping extender of the threading.Thread class.
class fitnessThread( threading.Thread ):
	def __init__(self,grid,offset,heightrange):
		threading.Thread.__init__(self)
		self.grid = grid
		self.offset = offset
		self.heightrange = heightrange
	def run(self):
		global fitness, imageDataArray, wrange
		myTotal = 0
		for h in self.heightrange:
			for w in wrange:
				cl = pg.Color(self.grid[h][w])
				r1,g1,b1 = cl[1],cl[2],cl[3]
				r2,g2,b2,a2 = imageDataArray[h][w]
				deltar = r1-r2
				deltag = g1-g2
				deltab = b1-b2
				myTotal += (deltar*deltar)+(deltag*deltag)+(deltab*deltab)
		fitness += myTotal

parts = 1
fitness = 0
numThreads = 2
imageDataArray = None

#global state hack for speed
wrange,hrange = None,None
def fitness_function(surface_1):
	global wrange, hrange, fitness, numThreads, parts

	fitness = 0
	pixelArray = pg.PixelArray(surface_1)

	threads = []

	#print 'numThreads:',numThreads

	for i in range(0,numThreads):
		start = i * parts
		end = (i+1)*parts
		# Create new threads
		threads.append(fitnessThread(pixelArray[start:end], start, range(0,parts)))

	#print 'starting threads'

	for t in threads:
		t.start()

	#print threading.enumerate()

	for t in threads:
		while t.isAlive():
			t.join()

	pixelArray = None

	return fitness


def main():
	if len(sys.argv) < 2:
		print "in order to use type ./genetic.py \"imagename\" where imagename is the image you want to approximate, followed by any arguments"
		return

	#some 'settings'
	numMutations = 1
	numPolygons = 15
	numSides = 4

	#init the pygame, grab the image and load it
	pg.init()
	imagename = sys.argv[1]
	imagesurf = pg.image.load(imagename)

	imsize = imagesurf.get_size()
	newsurf = pg.Surface(imsize)
	newsurf.fill((0,0,0))
	newsurf.blit(imagesurf,(0,0))
	imagesurf = newsurf

	#generate the target image array
	global imageDataArray
	imageDataArray = [0]*imsize[0]
	for i in range(0,imsize[0]):
		col = [0]*imsize[1]
		for j in range(0,imsize[1]):
			col[j] = imagesurf.get_at((i,j))
		imageDataArray[i] = col

	#create a drawing, initialise display output
	d = Drawing(imsize[0],imsize[1])
	d.Generate(numPolygons,numSides)
	width,height = imsize[0]*6,imsize[1]*2
	pg.display.set_mode((width,height))
	screen = pg.display.get_surface()
	screen.fill((0,0,0))
	screen.blit(imagesurf,(0,0))
	pg.display.update()

	#move d to parent for the upcoming loop, create a child
	parent = d
	child = d.Clone()

	#init the two drawing surfaces
	childSurface = pg.Surface(imsize)
	parentSurface = pg.Surface(imsize)
	drawpos = imsize[0],0
	drawpos2 = imsize[0]*2,0

	#generate wrange and hrange
	x,y = childSurface.get_size()
	global wrange,hrange
	wrange,hrange = range(0,x),range(0,y)

	#figure out the chopping up of the image data accross threads
	global numThreads, parts
	raw_parts = 1.0*y/numThreads
	parts = int(floor(raw_parts))
	if parts == 0:
		print 'you cannot have more threads than rows!'
		return
	if parts != raw_parts:
		numThreads += 1

	#get the target rectanlges for the child and parent - where they will get drawn
	parentRect = parentSurface.get_rect(topleft=drawpos)
	childRect = childSurface.get_rect(topleft=drawpos2)

	#draw the parent and find out how well it matches (badly at this point)
	parent.Draw(parentSurface)
	p = fitness_function(parentSurface)
	newParent = True
	renders = 0
	tnow = time()

	#for i in range(0,1000):
	while 1:
		#mutate and test the child
		child.Mutate(numMutations)
		child.Draw(childSurface)
		c = fitness_function(childSurface)

		assert parent != child
		#combine the surfaces for display
		if newParent:
			#only refresh the parent if its changed
			screen.fill((0,0,0),parentRect)
			screen.blit(parentSurface,drawpos)

		screen.fill((0,0,0),childRect)
		screen.blit(childSurface,drawpos2)
		pg.display.update()

		if c < p:
			#if the child is a better match
			#store the new 'best' match value
			p = c
			#move the child up to being the new parent
			parent = child.Clone()
			parentSurface.fill((0,0,0))
			parent.Draw(parentSurface)
			print p
			newParent = True
		else:
			#child is no good - begin afresh
			child = parent.Clone()
			newParent = False

		renders += 1
		if renders % 10 == 0:
			t = time()
			pg.display.set_caption(str(renders/(t-tnow)))
			tnow = t
			renders = 0

		if p == 0:
			raw_input('Done! (press enter to close)')
			break


if __name__ == "__main__":
	profiling = 0
	if profiling == 1:
		import cProfile,pstats
		prof = cProfile.Profile()
		prof = prof.runctx("main()", globals(), locals())
		stats = pstats.Stats(prof)
		stats.sort_stats("time")  # Or cumulative
		stats.print_stats(800)  # 80 = how many to print
		# The rest is optional.
		stats.print_callees()
		stats.print_callers()
	else: main()

